import os
import datetime
import gevent

from gevent.fileobject import FileObjectThread
from gevent.threadpool import ThreadPool

import websocket

from utils.file_md5 import file_md5

from app_config import APP_USER_CONFIG

from app.utils.ws_data import *
from ws_state import WSState

from .fod.file_sync_fod import Config as FileSyncFODConfig, FileSyncFOD

from app_instance import app

loop = gevent.get_hub().loop

class WSClient:
    '''websocket client'''
    def __init__(
            self,
            client_id=APP_USER_CONFIG['clientId'],
            ws_server=APP_USER_CONFIG['remoteServer'],
            data_dir=APP_USER_CONFIG['dataDir'],
            inbox_dir=APP_USER_CONFIG['inbox'],
            outbox_dir=APP_USER_CONFIG['outbox'],
            task_timeout_secs=APP_USER_CONFIG['taskTimeoutSeconds'],
            chunk_size = APP_USER_CONFIG['chunkSize']):
        enable_fod = False
        if 'enableFOD' in APP_USER_CONFIG:
            enable_fod = APP_USER_CONFIG['enableFOD']

        file_sync_fod_cfg = FileSyncFODConfig()
        if 'partSize' in APP_USER_CONFIG:
            file_sync_fod_cfg.part_size = APP_USER_CONFIG['partSize']
        if 'queueSize' in APP_USER_CONFIG:
            file_sync_fod_cfg.queue_size = APP_USER_CONFIG['queueSize']
        if 'costPerPart' in APP_USER_CONFIG:
            file_sync_fod_cfg.cost_per_part = APP_USER_CONFIG['costPerPart']
        if 'minCostTime' in APP_USER_CONFIG:
            file_sync_fod_cfg.min_cost_time = APP_USER_CONFIG['minCostTime']
        if 'maxCostTime' in APP_USER_CONFIG:
            file_sync_fod_cfg.max_cost_time = APP_USER_CONFIG['maxCostTime']
        if 'tmpPartInbox' in APP_USER_CONFIG:
            file_sync_fod_cfg.tmp_part_inbox = APP_USER_CONFIG['tmpPartInbox']
        file_sync_fod_cfg.part_inbox = inbox_dir
        file_sync_fod_cfg.part_outbox = outbox_dir
        file_sync_fod_cfg.data_dir = data_dir
        file_sync_fod_cfg.inbox = os.path.join(data_dir, '_inbox')
        file_sync_fod_cfg.outbox = os.path.join(data_dir, '_outbox')
        if enable_fod:
            self.file_sync_fod = FileSyncFOD(cfg=file_sync_fod_cfg, db_engine=app.db.engine)

        self.ws_server = ws_server

        self.client_id = client_id
        if self.client_id is None:
            self.client_id = APP_USER_CONFIG['clientId']

        self.chunk_size = chunk_size
        self.chunk_to_log = self.chunk_size * 10

        self.thread_pool = ThreadPool(4)

        self.has_pull_task = False
        self.has_push_task = False

        self.ws_stream_up = None
        self.ws_stream_down = None
        self.ws_task_pull = None
        self.ws_task_finish = None

        self.reconnect_state = True

        self.update_up_active_timestamp()
        self.update_down_active_timestamp()

        self.prefer_local = False

        self.wrap_gevent_proc(self.connect_to_server)

        self.state = WSState(
            self.client_id,
            loop,
            data_dir=data_dir,
            inbox_dir=file_sync_fod_cfg.inbox if enable_fod else inbox_dir,
            outbox_dir=file_sync_fod_cfg.outbox if enable_fod else outbox_dir,
            logger=app.logger,
            task_timeout_secs=task_timeout_secs,
            is_server=False,
            chunk_size=chunk_size)

        self.handle_task_timer = loop.timer(0.0, 0.1)
        self.handle_task_timer.start(self.handle_task_timer_handler)

        self.pull_task_timer = loop.timer(0.0, 3)
        self.pull_task_timer.start(self.wrap_gevent_proc, self.pull_task_timer_handler)

        self.check_state_timer = loop.timer(60, 60)
        self.check_state_timer.start(self.check_active_state)

    def update_up_active_timestamp(self):
        self.up_active_timestamp = datetime.datetime.now()

    def update_down_active_timestamp(self):
        self.down_active_timestamp = datetime.datetime.now()

    def check_active_state(self):
        delta = datetime.timedelta(seconds=30)
        now = datetime.datetime.now()
        if (self.has_push_task and now - self.up_active_timestamp > delta) or \
                (self.has_pull_task and now - self.down_active_timestamp > delta):
            app.logger.warning('upstream or downstream timeout, reconnecting...')
            self.reconnect_to_server()

    def connect_to_server(self):
        if self.reconnect_state == True:
            self.reconnect_state = False

            wss = [
                self.ws_task_pull,
                self.ws_task_finish,
                self.ws_stream_up,
                self.ws_stream_down
            ]

            for ws in wss:
                if ws:
                    ws.close()

            self.ws_task_pull = websocket.create_connection('ws://{0}/websocket/task/pull'.format(self.ws_server))
            self.ws_task_finish = websocket.create_connection('ws://{0}/websocket/task/finish'.format(self.ws_server))

            self.ws_stream_up = websocket.create_connection('ws://{0}/websocket/stream/up'.format(self.ws_server))
            self.has_push_task = False

            self.ws_stream_down = websocket.create_connection('ws://{0}/websocket/stream/down'.format(self.ws_server))
            self.has_pull_task = False

            app.logger.info('websocket server connected')

    def reconnect_to_server(self):
        self.reconnect_state = True
        self.spawn_later(3.0, self.connect_to_server)

    def wrap_gevent_proc(self, proc, *args, **kargs):
        try:
            proc(*args, **kargs)
        except Exception as ex:
            app.logger.error(ex)
            self.reconnect_to_server()

    def spawn(self, proc, *args, **kargs):
        gevent.spawn(self.wrap_gevent_proc, proc, *args, **kargs)

    def spawn_later(self, delay, proc, *args, **kargs):
        gevent.spawn_later(delay, self.wrap_gevent_proc, proc, *args, **kargs)

    def run_chunk_perf(self):
        ws_chunk_perf = websocket.create_connection('ws://{0}/websocket/perf/chunk/push'.format(self.ws_server))
        data = bytearray().zfill(5 * 1000 * 1024)
        for i in range(0, 100000):
            ws_chunk_perf.send_binary(data)
            ws_chunk_perf.recv()
            print(i)

    def recv_task(self):
        data = self.ws_task_pull.recv()
        if data and isinstance(data, bytes):
            opcode, header, _ = WSDataPack.unpack(data)
            if opcode == WSOPCode.ACK.value:
                if 'ack' in header and 'tasks' in header:
                    tasks = header['tasks']
                    for task in tasks:
                        assert 'opcode' in task
                        app.logger.debug('adding transient task %s %s' % (str(WSOPCode(task['opcode'])), task['file_name']))
                        self.state.add_transient_task(task)

        self.spawn_later(3.0, self.pull_task_timer_handler)

    def pull_task_timer_handler(self):
        if self.state.transient_task_count() > 0:
            return

        if self.ws_task_pull and not self.ws_task_pull.connected:
            raise ConnectionResetError('websocket is not connected')
        self.ws_task_pull.send_binary(
            WSDataPack.pack(WSOPCode.PULL_TASK.value, {'client_id': self.client_id}, None))

        self.spawn(self.recv_task)

    def handle_task_timer_handler(self):
        task = self.state.fetch_task(self.client_id, prefer_local=self.prefer_local)
        self.prefer_local = not self.prefer_local
        if task is not None:
            self.handle_task(task)

    def handle_task(self, task):
        if 'opcode' in task and 'file_name' in task:
            opcode = task['opcode']
            file_name = task['file_name']
            if opcode == WSOPCode.PULL_FILE.value:
                if self.state.is_finished_task(file_name):
                    self.state.finish_unfinished_state(file_name)
                    self.spawn(self.send_finish, file_name)
                elif self.state.is_unfinished_file(file_name):
                    task = self.state.fetch_unfinished_file_task(file_name)
                    assert task is not None
                    if self.state.is_task_timeout(task):
                        self.state.add_transient_task(task)
                        app.logger.debug('add unfinished timeout task %s' % file_name)
                    else:
                        app.logger.debug('ignore %s %s' % (str(WSOPCode(opcode)), file_name))
                else:
                    if self.has_pull_task:
                        self.state.return_task(task)
                    else:
                        app.logger.debug('handling WSOPCode.PULL_FILE %s' % file_name)
                        self.has_pull_task = True
                        self.spawn(self.pull_file, file_name, task['size'], task['checksum'])
            elif opcode == WSOPCode.PULL_CHUNKS.value:
                if self.state.is_finished_task(file_name):
                    self.state.finish_unfinished_state(file_name)
                    self.spawn(self.send_finish, file_name)
                elif self.has_pull_task:
                    if not self.state.is_unfinished_file(file_name):
                        self.state.return_task(task)
                else:
                    app.logger.debug('handling WSOPCode.PULL_CHUNKS %s' % file_name)
                    self.has_pull_task = True
                    self.spawn(self.pull_chunks, file_name, task['chunks'], task['size'], task['checksum'])
            elif opcode == WSOPCode.PUSH_FILE.value:
                if self.state.is_finished_file(file_name):
                    app.logger.warn('file task %s is finished' % file_name)
                else:
                    if self.has_push_task:
                        self.state.return_task(task)
                    else:
                        app.logger.debug('handling WSOPCode.PUSH_FILE %s' % file_name)
                        self.has_push_task = True
                        self.spawn(self.push_file, file_name, task['size'], task['checksum'])
            elif opcode == WSOPCode.PUSH_CHUNKS.value:
                if self.state.is_finished_file(file_name):
                    app.logger.warn('file task %s is finished' % file_name)
                else:
                    if self.has_push_task:
                        self.state.return_task(task)
                    else:
                        self.state.remove_push_task(file_name)
                        app.logger.debug('handling WSOPCode.PUSH_CHUNKS %s' % file_name)
                        self.has_push_task = True
                        self.spawn(self.push_chunks, file_name, task['chunks'])
            elif opcode == WSOPCode.FINISH_TASK.value:
                app.logger.debug('task %s finished' % file_name)
                self.state.finish_task(file_name, self.client_id)
        else:
            app.logger.warning('opcode %s not in task')

    def write_chunk(self, file_name, fobject, offset, chunk_size, close_fd=True):
        if fobject.closed:
            return (True, offset)
        fobject.seek(offset)
        chunk = fobject.read(chunk_size)
        if chunk is not None and len(chunk) > 0:
            self.ws_stream_up.send_binary(WSDataPack.pack(
                WSOPCode.PUSH_CHUNK.value,
                {
                    'file_name': file_name,
                    'offset': offset,
                },
                chunk))

            self.update_up_active_timestamp()

            sent_size = offset + len(chunk)
            if sent_size % self.chunk_to_log == 0:
                app.logger.debug('sent %sM(%s)' % (sent_size / 1024000, file_name))

            data = self.ws_stream_up.recv()
            if not isinstance(data, bytes):
                app.logger.debug('invalid ack data for push trunks')
                self.has_push_task = False
                return

            opcode, header, _ = WSDataPack.unpack(data)
            if opcode != WSOPCode.ACK.value or \
                'ack' not in header or \
                header['ack'] != WSOPCode.PUSH_CHUNK.value:
                return

            self.state.update_unack_task_modify_t(file_name)
            return (False, offset + len(chunk))
        else:
            if close_fd:
                app.logger.debug('close file %s' % file_name)
                fobject.close()
            self.has_push_task = False
            return (True, offset)

    def recv_push_file_ack(self, file_name, file_path):
        data = self.ws_stream_up.recv()
        if not isinstance(data, bytes):
            app.logger.debug('invalid ack data for push file')
            self.has_push_task = False
            return

        opcode, header, _ = WSDataPack.unpack(data)
        if opcode != WSOPCode.ACK.value or 'ack' not in header:
            app.logger.debug('invalid ack data for push file')
            self.has_push_task = False
            return

        ack_opcode = header['ack']
        if  ack_opcode  != WSOPCode.PUSH_FILE.value:
            app.logger.warn('ack error %s' % str(WSOPCode(ack_opcode)))
            self.has_push_task = False
            return

        app.logger.debug('opening file %s in rb mode' % file_name)
        fobject = FileObjectThread(open(file_path, 'rb'), threadpool=self.thread_pool)

        offset = 0
        while True:
            r, offset = self.write_chunk(file_name, fobject, offset, self.chunk_size)
            if r:
                app.logger.debug('write file %s finished' % file_name)
                break
            gevent.sleep(0.1)

    def push_file(self, file_name, size, checksum):
        file_path = os.path.join(self.state.outbox, file_name)
        if not os.path.exists(file_path):
            app.logger.warn('%s is not exists' % file_path)
            self.has_push_task = False
            return

        self.ws_stream_up.send_binary(
            WSDataPack.pack(
                WSOPCode.PUSH_FILE.value,
                {
                    'client_id': self.client_id,
                    'file_name': file_name,
                    'size': size,
                    'checksum': checksum
                }
            ))

        app.logger.debug('request sent, file name: %s' % file_name)
        self.spawn(self.recv_push_file_ack, file_name, file_path)

    def recv_push_chunks_ack(self, file_name, file_path, chunks):
        data = self.ws_stream_up.recv()
        if not isinstance(data, bytes):
            app.logger.error('invalid ack data for WSOPCode.PUSH_FILE')
            self.has_push_task = False
            return

        opcode, header, _ = WSDataPack.unpack(data)

        opcode, header, _ = WSDataPack.unpack(data)
        if opcode != WSOPCode.ACK.value or 'ack' not in header:
            app.logger.debug('invalid ack data for WSOPCode.PUSH_FILE')
            self.has_push_task = False
            return

        ack_opcode = header['ack']
        if ack_opcode != WSOPCode.PUSH_FILE.value or ack_opcode == WSOPCode.PUSH_CHUNKS.value:
            app.logger.warn('ack error %s' % str(WSOPCode(ack_opcode)))
            self.has_push_task = False
            return

        app.logger.debug('opening file %s in rb mode' % file_name)
        fobject = FileObjectThread(open(file_path, 'rb'), threadpool=self.thread_pool)
        for chunk_item in chunks:
            if len(chunk_item) != 2:
                continue
            offset = chunk_item[0]
            size = offset + chunk_item[1]
            while offset < size:
                r, offset = self.write_chunk(file_name, fobject, offset, self.chunk_size, close_fd=False)
                if r:
                    app.logger.debug('write file %s finished' % file_name)
                    break
                gevent.sleep(0.1)

        fobject.close()
        app.logger.debug('close file %s' % file_name)

        self.has_push_task = False

    def push_chunks(self, file_name, chunks=[]):
        file_path = os.path.join(self.state.outbox, file_name)
        if not os.path.exists(file_path):
            app.logger.warn('%s is not exists' % file_path)
            self.has_push_task = False
            return
        self.ws_stream_up.send_binary(
            WSDataPack.pack(
                WSOPCode.PUSH_CHUNKS.value,
                {
                    'file_name': file_name,
                    'client_id': self.client_id,
                    'size': os.path.getsize(file_path),
                    'checksum': file_md5(file_path)
                },
                None))
        self.spawn(self.recv_push_chunks_ack, file_name, file_path, chunks)

    def recv_chunk(self, file_name, fobject, size, checksum):
        app.logger.debug('receiving %s chunk' % file_name)

        data = self.ws_stream_down.recv()
        if not data or not isinstance(data, bytes):
            app.logger.error('invalid data in recv_chunk')
            return
        opcode, header, chunk = WSDataPack.unpack(data)
        if 'offset' not in header:
            app.logger.error('invalid key "offset" in header')
            return

        app.logger.debug('received %s chunk' % file_name)

        offset = header['offset']
        fobject.seek(offset)
        fobject.write(chunk)
        fobject.flush()

        self.update_down_active_timestamp()

        received_size = offset + len(chunk)
        if received_size % self.chunk_to_log == 0:
            app.logger.debug('received %sM(%s)' % (received_size / 1024000, file_name))

        self.ws_stream_down.send_binary(WSDataPack.pack(WSOPCode.ACK.value, {'ack': WSOPCode.PUSH_CHUNK.value}))

        app.logger.debug('send %s chunk ack' % file_name)

        unfinished_state = self.state.get_unfinished_state(file_name)

        assert unfinished_state is not None

        if self.state.update_unfinished_state(file_name, (offset, len(chunk))):
            app.logger.debug('finished %s' % file_name)

            self.has_pull_task = False
            fobject.close()

            app.logger.debug('file %s closed' % file_name)
            tmp_file_path = self.state.get_tmp_file_path(file_name)

            try:
                verify_succeed = True

                app.logger.debug('verify file %s\'s checksum' % tmp_file_path)
                md5 = file_md5(tmp_file_path)
                if md5 != checksum:
                    verify_succeed = False
                    app.logger.error('file %s\'s checksum not match, expect %s, actual %s' % (tmp_file_path, checksum, md5))

                if verify_succeed:
                    self.state.finish_task(file_name, self.client_id)
                    self.spawn(self.send_finish, file_name)
                    self.state.add_finish_task_record(unfinished_state)

                app.logger.debug('finishing task %s' % file_name)
                self.state.finish_unfinished_state(file_name)

            except Exception as ex:
                app.logger.error('exception %s happend while checking file %s\'s checksum' % (ex, file_name))

            return

        self.spawn(self.recv_chunk, file_name, fobject, size, checksum)

    def recv_file(self, file_name, size, checksum):
        if self.state.is_finished_task(file_name):
            self.has_pull_task = False
            return

        data = self.ws_stream_down.recv()
        if not data or not isinstance(data, bytes):
            return
        opcode, header, _ = WSDataPack.unpack(data)
        if opcode != WSOPCode.ACK.value or 'ack' not in header:
            self.has_pull_task = False
            app.logger.debug('invalid task %s' % file_name)
            return
        file_path = self.state.get_tmp_file_path(file_name)

        r = self.state.add_unfinished_state(
            {'file_name': file_name, 'size': size, 'clients_id': [self.client_id], 'checksum': checksum})

        app.logger.debug('opening file %s in wb+ mode' % file_name)
        fobject = FileObjectThread(open(file_path, 'wb+' if r else 'rb+'), threadpool=self.thread_pool)
        app.logger.debug('opened file %s in wb+ mode' % file_name)

        if size == 0:
            self.state.finish_unfinished_state(file_name)
            self.spawn(self.send_finish, file_name)
            fobject.close()
            self.has_pull_task = False
            app.logger.debug('file %s closed' % file_name)
        else:
            self.spawn(self.recv_chunk, file_name, fobject, size, checksum)

    def send_finish(self, file_name):
        app.logger.debug('adding finish task, file name: %s' % file_name)
        self.ws_task_finish.send_binary(
            WSDataPack.pack(WSOPCode.FINISH_TASK.value,
            {'file_name': file_name, 'client_id': self.client_id}))

    def pull_file(self, file_name, size, checksum):
        if size == 0:
            self.state.finish_unfinished_state(file_name)
            self.spawn(self.finish_file, file_name)
            return
        app.logger.debug('pulling file %s' % file_name)
        self.ws_stream_down.send_binary(
            WSDataPack.pack(
                WSOPCode.PULL_FILE.value,
                {
                    'client_id': self.client_id,
                    'file_name': file_name,
                }))
        app.logger.debug('receiving file %s' % file_name)
        self.spawn(self.recv_file, file_name, size, checksum)

    def recv_chunks(self, file_name, size, checksum):
        if self.state.is_finished_task(file_name):
            self.has_pull_task = False
            return

        data = self.ws_stream_down.recv()
        if not data or not isinstance(data, bytes):
            self.has_pull_task = False
            return
        opcode, header, _ = WSDataPack.unpack(data)

        if opcode != WSOPCode.ACK.value or 'ack' not in header:
            app.logger.debug('invalid ack data for pull trunks')
            self.has_push_task = False
            return

        ack_opcode = header['ack']

        if ack_opcode == WSOPCode.PULL_CHUNKS.value:
            file_path = self.state.get_tmp_file_path(file_name)
            if os.path.exists(file_path):
                app.logger.debug('opening file %s in rb+ mode' % file_name)
                fobject = FileObjectThread(open(file_path, 'rb+'), threadpool=self.thread_pool)
                self.spawn(self.recv_chunk, file_name, fobject, size, checksum)
            else:
                app.logger.error('file %s is not exists while receive chunks' % file_path)
        else:
            if ack_opcode == WSOPCode.TASK_FINISHED.value or \
                ack_opcode == WSOPCode.TASK_NOT_EXIST.value:
                self.state.finish_unfinished_state(file_name)
            app.logger.warn('ack error %s' % str(WSOPCode(ack_opcode)))
            self.has_pull_task = False

    def pull_chunks(self, file_name, chunks, size, checksum):
        if size == 0:
            self.state.update_unfinished_state(file_name, (0, 0))
            self.spawn(self.finish_file, file_name)
            return
        self.ws_stream_down.send_binary(
            WSDataPack.pack(
                WSOPCode.PULL_CHUNKS.value,
                {
                    'client_id': self.client_id,
                    'file_name': file_name,
                    'chunks': chunks
                }))
        self.spawn(self.recv_chunks, file_name, size, checksum)

wc_client = WSClient()

if __name__ == "__main__":
    #websocket.enableTrace(True)
    from gevent import pywsgi

    server = pywsgi.WSGIServer(('0.0.0.0', 8001), app)
    server.serve_forever()
